﻿// dllmain.cpp : 定义 DLL 应用程序的入口点。
#include "pch.h"
#include <windows.h>
#include <winsock2.h>
#include <iostream>
#include <fstream>

#pragma comment(lib, "ws2_32.lib")

using namespace std;

BYTE hook[6];
BYTE hook2[6];
BYTE jmp[6] = { 0xe9,0x00, 0x00, 0x00, 0x00 ,0xc3 };
ofstream myfile;
ofstream myfile2;
DWORD pPrevious;

DWORD HookFunction(LPCSTR lpModule, LPCSTR lpFuncName, LPVOID lpFunction, unsigned char* lpBackup)
{
    DWORD dwAddr = (DWORD)GetProcAddress(GetModuleHandle(lpModule), lpFuncName);
    ReadProcessMemory(GetCurrentProcess(), (LPVOID)dwAddr, lpBackup, 6, 0);
    DWORD dwCalc = ((DWORD)lpFunction - dwAddr - 5);
    VirtualProtect((void*)dwAddr, 6, PAGE_EXECUTE_READWRITE, &pPrevious);
    memcpy(&jmp[1], &dwCalc, 4);
    WriteProcessMemory(GetCurrentProcess(), (LPVOID)dwAddr, jmp, 6, 0);
    VirtualProtect((void*)dwAddr, 6, pPrevious, &pPrevious);
    FlushInstructionCache(GetCurrentProcess(), 0, 0);
    return dwAddr;
}

BOOL UnHookFunction(LPCSTR lpModule, LPCSTR lpFuncName, unsigned char* lpBackup)
{
    DWORD dwAddr = (DWORD)GetProcAddress(GetModuleHandle(lpModule), lpFuncName);

    if (WriteProcessMemory(GetCurrentProcess(), (LPVOID)dwAddr, lpBackup, 6, 0))
        return TRUE;
    FlushInstructionCache(GetCurrentProcess(), 0, 0);

    return FALSE;
}

int __stdcall nSend(SOCKET s, const char* buf, int len, int flags) {
    UnHookFunction("ws2_32.dll", "send", hook);
    int result = send(s, buf, len, flags);
    
    myfile.open("D:\\sendlog.txt", ios::app | ios::binary);
    myfile << buf;
    myfile.close();


    HookFunction("ws2_32.dll", "send", (LPVOID*)nSend, hook);
    return result;
}

int __stdcall nRecv(SOCKET s, char* buf, int len, int flags)
{
    UnHookFunction("ws2_32.dll", "recv", hook2);
    len = recv(s, buf, len, flags);

    if (len > 0)
    {
        myfile2.open("D:\\recvlog.txt", ios::app | ios::binary);
        myfile2 << buf;
        myfile2.close();
    }
    HookFunction("ws2_32.dll", "recv", (LPVOID*)nRecv, hook2);
    return len;
}

void regist() {
    HookFunction("ws2_32.dll", "send", (LPVOID*)nSend, hook);
    HookFunction("ws2_32.dll", "recv", (LPVOID*)nRecv, hook2);
}

void unregist() {
    UnHookFunction("ws2_32.dll", "send", hook);
    UnHookFunction("ws2_32.dll", "recv", hook2);
}

BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserved)
{
    switch (ul_reason_for_call)
    {
    case DLL_PROCESS_ATTACH:
        regist();
    case DLL_THREAD_ATTACH:
    case DLL_THREAD_DETACH:
        break;
    case DLL_PROCESS_DETACH:
        unregist();
        break;
    }
    return TRUE;
}